import prometheus_client
from prometheus_client import Counter, Gauge
from prometheus_client.core import CollectorRegistry
from sqlalchemy.orm import Session
from datetime import datetime, timedelta

from core.logger import logger
from common.utils import CommonTimeUtils
from app.monitor_metrics.models.prom_alarm_record import PromAlarmRecord
from app.celery_task.models.celery_task import CeleryTask


class PrometheusMetricsHandler(object):

    def __init__(self):

        self.registry = CollectorRegistry(auto_describe=False)
        self.celery_state = Gauge(name="celery_state",
                           documentation="Celery is running. 1:ok, 0:failed",
                           registry=self.registry)
        self.alarm_state = Gauge(name="alarm_state",
                             documentation="prometheus alarm is running. 1:ok, 0:failed",
                              registry=self.registry)
        
        # 带标签的用法
        # self.mq_messages = Gauge("mq_messages",
        #                      "Number of queue messages.",
        #                      ["cluster", "queue"], registry=self.registry)
        # self.mq_messages.labels(cluster=cluster, queue=queue_name).set(queue_info["messages_ready"])


    def get_all_metrics(self, db: Session):
        # 设置celery 工作状态
        self.set_celery_state(db)
        
        # 设置prometheus告警 工作状态
        self.set_alarm_state(db)
        
        return prometheus_client.generate_latest(self.registry)


    def set_celery_state(self, db: Session):
        """
        设置celery 工作状态
        """
        try:
            # 
            celery_watchdog = CeleryTask.get_object_by_name(db, "watch_dog")
            if not celery_watchdog:
                logger.error(f"celery watchdog not exist")
                self.celery_state.set(0)
                return False

            celery_watchdog_info = celery_watchdog.object_to_dict(db)
            now_timestamp = datetime.now().timestamp()
            update_timestamp = datetime.timestamp(celery_watchdog_info["update_at"] + timedelta(minutes=15))
            if update_timestamp < now_timestamp:
                logger.error(f"celery watchdog no running last 15 minutes")
                self.celery_state.set(0)
                return False
        except Exception as e:
            logger.exception(f"set celery state error => {e}")
            self.celery_state.set(0)
            return False

        self.celery_state.set(1)
        return True
    
    
    def set_alarm_state(self, db: Session):
        """
        设置prometheus告警 工作状态
        """
        try:
            # 检查watchdog的告警记录是否存在
            watchdog_info = PromAlarmRecord.get_object_info_for_watchdog(db)
            if not watchdog_info:
                logger.error(f"watchdog not exist")
                self.alarm_state.set(0)
                return False

            # 检查watchdog的告警记录是否超过15分钟
            firing_infos = PromAlarmRecord.get_object_info_for_firing(db)
            now_timestamp = datetime.now().timestamp()
            for firing_info in firing_infos:
                firing_timestamp = datetime.timestamp(CommonTimeUtils.strptime(firing_info["update_at"]) + timedelta(minutes=15))
                if firing_timestamp < now_timestamp:
                    logger.error(f"no alarm last 15 minutes => {firing_info}")
                    self.alarm_state.set(0)
                    return False
        except Exception as e:
            logger.exception(f"set alarm state error => {e}")
            self.alarm_state.set(0)
            return False

        self.alarm_state.set(1)
        return True


def main():

    prometheus_metrics = PrometheusMetrics()
    queue_data = prometheus_metrics.get_all_metrics()
    pprint(queue_data)
    #pprint(METRICS_QUEUE)

if __name__ == "__main__":
    main()
